package demo.runtime.dispatcher;

import org.apache.flink.api.common.JobID;
import org.apache.flink.api.common.time.Time;
import org.apache.flink.runtime.client.JobExecutionException;
import org.apache.flink.runtime.client.JobSubmissionException;
import org.apache.flink.runtime.concurrent.FutureUtils;
import org.apache.flink.runtime.dispatcher.DispatcherException;
import org.apache.flink.runtime.dispatcher.DispatcherId;
import org.apache.flink.runtime.dispatcher.DispatcherServices;
import org.apache.flink.runtime.dispatcher.JobManagerRunnerFactory;
import org.apache.flink.runtime.executiongraph.ArchivedExecutionGraph;
import org.apache.flink.runtime.jobgraph.JobGraph;
import org.apache.flink.runtime.jobmanager.JobGraphWriter;
import org.apache.flink.runtime.jobmaster.JobManagerRunner;
import org.apache.flink.runtime.jobmaster.JobNotFinishedException;
import org.apache.flink.runtime.jobmaster.factories.DefaultJobManagerJobMetricGroupFactory;
import org.apache.flink.runtime.messages.Acknowledge;
import org.apache.flink.runtime.rpc.FencedRpcEndpoint;
import org.apache.flink.runtime.rpc.PermanentlyFencedRpcEndpoint;
import org.apache.flink.runtime.rpc.RpcService;
import org.apache.flink.util.ExceptionUtils;
import org.apache.flink.util.Preconditions;
import org.apache.flink.util.function.BiConsumerWithException;
import org.apache.flink.util.function.FunctionUtils;
import org.apache.flink.util.function.FunctionWithException;

import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;

public class DispatcherDemo extends PermanentlyFencedRpcEndpoint<DispatcherId> {

    private final Map<JobID, CompletableFuture<JobManagerRunner>> jobManagerRunnerFutures = new HashMap<>(16);
    private final Map<JobID, CompletableFuture<Void>> jobManagerTerminationFutures = new HashMap<>(2);;
    private final JobManagerRunnerFactory jobManagerRunnerFactory;

    private final JobGraphWriter jobGraphWriter;

    protected DispatcherDemo(
            RpcService rpcService,
            String endpointId,
            DispatcherId fencingToken,
            DispatcherServicesDemo dispatcherServices
    ) {
        super(rpcService, endpointId, fencingToken);
        this.jobGraphWriter = dispatcherServices.getJobGraphWriter();
        this.jobManagerRunnerFactory = dispatcherServices.getJobManagerRunnerFactory();
    }

    public CompletableFuture<Acknowledge> submitJob(JobGraph jobGraph, Time timeout) {
        return internalSubmitJob(jobGraph);
    }

    private CompletableFuture<Acknowledge> internalSubmitJob(JobGraph jobGraph) {
        final CompletableFuture<Acknowledge> persistAndRunFuture = waitForTerminatingJobManager(jobGraph.getJobID(), jobGraph, this::persistAndRunJob)
                .thenApply(ignored -> Acknowledge.get());

        return persistAndRunFuture.handleAsync((acknowledge, throwable) -> {
            if (throwable == null) {
                return acknowledge;
            }
        }, getRpcService().getExecutor());
    }

    private CompletableFuture<Void> persistAndRunJob(JobGraph jobGraph) throws Exception {
        jobGraphWriter.putJobGraph(jobGraph);

        final CompletableFuture<Void> runJobFuture = runJob(jobGraph);

        return runJobFuture.whenComplete(BiConsumerWithException.unchecked((Object ignored, Throwable throwable) -> {
            if (throwable != null) {
                jobGraphWriter.removeJobGraph(jobGraph.getJobID());
            }
        }));
    }

    private CompletableFuture<Void> runJob(JobGraph jobGraph) {
        Preconditions.checkState(!jobManagerRunnerFutures.containsKey(jobGraph.getJobID()));

        final CompletableFuture<JobManagerRunner> jobManagerRunnerFuture = createJobManagerRunner(jobGraph);

        jobManagerRunnerFutures.put(jobGraph.getJobID(), jobManagerRunnerFuture);

        return jobManagerRunnerFuture
                .thenApply(FunctionUtils.uncheckedFunction(this::startJobManagerRunner))
                .thenApply(FunctionUtils.nullFn())
                .whenCompleteAsync(
                        (ignored, throwable) -> {
                            if (throwable != null) {
                                jobManagerRunnerFutures.remove(jobGraph.getJobID());
                            }
                        },
                        getMainThreadExecutor());
    }

    private JobManagerRunner startJobManagerRunner(JobManagerRunner jobManagerRunner) throws Exception {
        final JobID jobId = jobManagerRunner.getJobID();

        FutureUtils.assertNoException(
                jobManagerRunner.getResultFuture().handleAsync(
                        (ArchivedExecutionGraph archivedExecutionGraph, Throwable throwable) -> {
                            // check if we are still the active JobManagerRunner by checking the identity
                            final JobManagerRunner currentJobManagerRunner = Optional.ofNullable(jobManagerRunnerFutures.get(jobId))
                                    .map(future -> future.getNow(null))
                                    .orElse(null);
                            //noinspection ObjectEquality
                            if (jobManagerRunner == currentJobManagerRunner) {
                                if (archivedExecutionGraph != null) {
                                    jobReachedGloballyTerminalState(archivedExecutionGraph);
                                } else {
                                    final Throwable strippedThrowable = ExceptionUtils.stripCompletionException(throwable);

                                    if (strippedThrowable instanceof JobNotFinishedException) {
                                       // jobNotFinished(jobId);
                                    } else {
                                      //  jobMasterFailed(jobId, strippedThrowable);
                                    }
                                }
                            } else {
                                log.debug("There is a newer JobManagerRunner for the job {}.", jobId);
                            }

                            return null;
                        }, getMainThreadExecutor()));

        jobManagerRunner.start();

        return jobManagerRunner;
    }

    private CompletableFuture<JobManagerRunner> createJobManagerRunner(JobGraph jobGraph) {
        final RpcService rpcService = getRpcService();

        return CompletableFuture.supplyAsync(
                () -> {
                    try {
                        return jobManagerRunnerFactory.createJobManagerRunner(
                                jobGraph,
                                configuration,
                                rpcService,
                                highAvailabilityServices,
                                heartbeatServices,
                                jobManagerSharedServices,
                                new DefaultJobManagerJobMetricGroupFactory(jobManagerMetricGroup),
                                fatalErrorHandler);
                    } catch (Exception e) {
                        throw new CompletionException(new JobExecutionException(jobGraph.getJobID(), "Could not instantiate JobManager.", e));
                    }
                },
                rpcService.getExecutor());
    }

    private CompletableFuture<Void> waitForTerminatingJobManager(JobID jobId, JobGraph jobGraph, FunctionWithException<JobGraph, CompletableFuture<Void>, ?> action) {
        final CompletableFuture<Void> jobManagerTerminationFuture = getJobTerminationFuture(jobId);

        return jobManagerTerminationFuture.thenComposeAsync(
                FunctionUtils.uncheckedFunction((ignored) -> {
                    jobManagerTerminationFutures.remove(jobId);
                    return action.apply(jobGraph);
                }),
                getMainThreadExecutor());
    }

    CompletableFuture<Void> getJobTerminationFuture(JobID jobId) {
        if (!jobManagerRunnerFutures.containsKey(jobId)) {
            return FutureUtils.completedExceptionally(new DispatcherException(String.format("Job with job id %s is still running.", jobId)));
        } else {
            return jobManagerTerminationFutures.getOrDefault(jobId, CompletableFuture.completedFuture(null));
        }

    }
}
